import numpy as np
import torch
from torch.distributions import Categorical


# Agent no communication
class Agent:
    def __init__(self, algorithm, state_dim, action_dim, args):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.algorithm = algorithm
        self.args = args
        if self.algorithm == 'DQN':
            from agent.dqn_agent import DQNAgent
            self.agent = DQNAgent(self.state_dim, self.action_dim, self.args, False)
        elif self.algorithm == 'DDQN':
            from agent.dqn_agent import DQNAgent
            self.agent = DQNAgent(self.state_dim, self.action_dim, self.args, True)
        elif self.algorithm == 'RainbowDQN':
            from agent.rainbow_dqn_agent import RainbowDQNAgent
            self.agent = RainbowDQNAgent(self.state_dim, self.action_dim, self.args, double_dqn=True, dueling_dqn=True, n_step_dqn=False, prioritized_buffer=True)
        elif self.algorithm == 'DDPG':
            from agent.ddpg_agent import DDPGAgent
            self.agent = DDPGAgent(self.state_dim, self.action_dim, self.args)
        else:
            raise Exception("No such algorithm")
        

    def choose_action(self, state, exploration=True):
        if self.algorithm == 'DQN':
            if exploration:
                return self.agent.e_greedy_action(state)
            else:
                return self.agent.action(state)
        elif self.algorithm == 'DDQN':
            if exploration:
                return self.agent.e_greedy_action(state)
            else:
                return self.agent.action(state)
        elif self.algorithm == 'RainbowDQN':
            if exploration:
                return self.agent.e_greedy_action(state)
            else:
                return self.agent.action(state)
        elif self.algorithm == 'DDPG':
            if exploration:
                return self.agent.random_selection_action(state)
            else:
                return self.agent.action(state)
        else:
            raise Exception("No such algorithm")

    def remember(self, state, action, reward, next_state, done ):
        self.agent.replay_buffer.remember(state,action,reward,next_state,done)

    def train(self,):
        return self.agent.learn()
    
    def replayBufferLen(self, ):
        return len(self.agent.replay_buffer)

    def save(self,episodes):
        self.agent.save(episodes)

    def load(self,path):
        self.agent.load(path)



